# -*- coding: utf-8 -*-#

# -------------------------------------------------------------------------------
# Name:         ma_strategy
# Description:  双均线策略
# Author:       ylf
# Date:         8/19/21
# -------------------------------------------------------------------------------

import pandas as pd
import data.stock as st
import numpy as np
import strategy.base as strb
import matplotlib.pyplot as plt


def ma_stragecy(data, short_windows=5, long_windows=20):
    """
    双均线策略
    :param data:传入的dataframe数据
    :param short_windows: 短期均线，默认为5
    :param long_windows: 长期均线，默认为20
    :return: 返回一个整合过的dataframe
    """
    print("当前参数为：", short_windows, long_windows)
    # 1  计算出5日平均线数据和20日平均线数据
    data = pd.DataFrame(data)
    data['ma5'] = data['close'].rolling(short_windows).mean()
    data['ma20'] = data['close'].rolling(long_windows).mean()
    # 2 计算出金叉和死叉的指标
    data['buy_signal'] = np.where(data['ma5'] > data['ma20'], 1, 0)
    data['sell_signal'] = np.where(data['ma5'] < data['ma20'], -1, 0)
    # 3 对金叉和死叉的指标进行整合,去掉连续的1和-1都变成0，去掉所有的0
    data = strb.conform_signal(data)
    data = strb.fix_signal(data)
    # 4  计算单次收益率和累计收益率
    data = strb.calculate_return_rate(data)
    data = strb.calculate_sum_rate(data)
    return data


if __name__ == '__main__':
    # df = st.get_csv_price('000607.XSHE', '2015-01-01', '2021-08-01')
    # df = ma_stragecy(df)
    # print('共计买卖次数：', int(len(df)))
    # print(df[['close', 'ma5', 'ma20', 'singal', 'per_income', 'sum_rate']])

    # 计算多只股票的累计收益率，并画图展示
    stocks = ['000607.XSHE', '000017.XSHE', '000858.XSHE']
    # 构建一个新的dataframe
    new_frame = pd.DataFrame()
    for stock_code in stocks:
        df = st.get_csv_price(stock_code, '2016-01-01', '2021-08-01')
        df = ma_stragecy(df)
        # 抛弃日期索引
        new_frame[stock_code] = df['sum_rate'].reset_index(drop=True)
    plt.plot(new_frame)
    plt.title('MA')
    plt.legend()
    plt.show()




